#include "ASTPrinter.h"
#include <iostream>
#include <fstream>
#include <vector>
#include <stack>
#include <map>
#include <string>
#include <algorithm>
using namespace std;

// 词法符号结构
struct Symbol
{
    string type; // 词法类别，比如 "int", "Ident", "=", "EOF"
    string text; // 原始文本
    string code;
};

// 读取token流（词法分析器输出）
vector<Symbol> loadTokens(const string &filename)
{
    vector<Symbol> tokens;
    ifstream fin(filename);
    if (!fin)
    {
        cerr << "无法打开文件: " << filename << endl;
        return tokens;
    }
    string word, tokenTypeValue;
    while (fin >> word >> tokenTypeValue)
    {
        size_t start = tokenTypeValue.find('<');
        size_t comma = tokenTypeValue.find(',');
        size_t end = tokenTypeValue.find('>');
        if (start == string::npos || comma == string::npos || end == string::npos)
            continue;
        string type = tokenTypeValue.substr(start + 1, comma - start - 1);
        string val = tokenTypeValue.substr(comma + 1, end - comma - 1);
        tokens.push_back({type, word, val});
    }
    fin.close();
    tokens.push_back({"EOF", "$", "$"}); // 添加终结符EOF方便分析结束判断
    return tokens;
}

// 打印当前分析动作
void printAction(int step, const string &stackTop, const Symbol &curToken, const string &action)
{
    cout << step << "\t" << stackTop << " # " << curToken.text << "\t" << action << endl;
}

string mapTokenTypeToTerminal(const Symbol &token)
{
    if (token.type == "KW") // 关键字
    {
        if (token.text == "int")
            return "int";
        if (token.text == "void")
            return "void";
        if (token.text == "return")
            return "return";
        if (token.text == "const")
            return "const";
        if (token.text == "if")
            return "if";
        if (token.text == "else")
            return "else";
        if (token.text == "float")
            return "float";
        if (token.text == "main")
            return "Ident";
    }
    else if (token.type == "IDN") // 标识符
    {
        return "Ident";
    }
    else if (token.type == "INT") // 整数常量
    {
        return "IntConst";
    }
    else if (token.type == "FLOAT") // 浮点常量
    {
        return "FloatConst";
    }
    else if (token.type == "OP") // 运算符
    {
        return token.text;
    }
    else if (token.type == "SE") // 分隔符
    {
        return token.text;
    }
    else if (token.type == "EOF")
    {
        return "EOF";
    }

    return token.type;
}

// 递归下降AST构建器
class RecursiveDescentASTBuilder
{
private:
    vector<Symbol> tokens;
    size_t tokenIndex;

public:
    RecursiveDescentASTBuilder(const vector<Symbol> &toks) : tokens(toks), tokenIndex(0) {}

    std::unique_ptr<ASTNode> buildAST()
    {
        tokenIndex = 0;
        return parseProgram();
    }

private:
    std::unique_ptr<ASTNode> parseProgram()
    {
        auto program = std::make_unique<ProgramNode>();
        auto compUnit = parseCompUnit();
        if (compUnit)
        {
            program->addChild(std::move(compUnit));
        }
        return std::move(program);
    }

    std::unique_ptr<ASTNode> parseCompUnit()
    {
        auto compUnit = std::make_unique<CompUnitNode>();

        while (tokenIndex < tokens.size() - 1)
        { // 跳过EOF
            auto item = parseTopLevelItem();
            if (item)
            {
                compUnit->addChild(std::move(item));
            }
            else
            {
                break;
            }
        }

        return std::move(compUnit);
    }

    std::unique_ptr<ASTNode> parseTopLevelItem()
    {
        if (tokenIndex >= tokens.size())
            return nullptr;

        if (tokens[tokenIndex].type == "KW" && tokens[tokenIndex].text == "const")
        {
            return parseConstDecl();
        }
        else if (isVariableDeclaration())
        {
            return parseVarDecl();
        }
        else if (isFunctionDefinition())
        {
            return parseFuncDef();
        }

        return nullptr;
    }

    bool isVariableDeclaration()
    {
        if (tokenIndex >= tokens.size())
            return false;

        size_t lookAhead = tokenIndex;

        // 跳过类型
        if (tokens[lookAhead].type == "KW" &&
            (tokens[lookAhead].text == "int" || tokens[lookAhead].text == "float"))
        {
            lookAhead++;
        }
        else
        {
            return false;
        }

        // 检查标识符
        if (lookAhead < tokens.size() && tokens[lookAhead].type == "IDN")
        {
            lookAhead++;
            // 检查后面不是左括号（不是函数）
            return lookAhead >= tokens.size() ||
                   mapTokenTypeToTerminal(tokens[lookAhead]) != "(";
        }

        return false;
    }

    bool isFunctionDefinition()
    {
        if (tokenIndex >= tokens.size())
            return false;

        size_t lookAhead = tokenIndex;

        // 跳过返回类型
        if (tokens[lookAhead].type == "KW" &&
            (tokens[lookAhead].text == "int" ||
             tokens[lookAhead].text == "float" ||
             tokens[lookAhead].text == "void"))
        {
            lookAhead++;
        }
        else
        {
            return false;
        }

        // 检查标识符
        if (lookAhead < tokens.size() &&
            (tokens[lookAhead].type == "IDN" || tokens[lookAhead].text == "main"))
        {
            lookAhead++;
            // 检查左括号
            return lookAhead < tokens.size() &&
                   mapTokenTypeToTerminal(tokens[lookAhead]) == "(";
        }

        return false;
    }

    std::unique_ptr<ASTNode> parseConstDecl()
    {
        // const float b1 = 8.2;
        auto constDecl = std::make_unique<ConstDeclNode>();

        // 跳过 const
        tokenIndex++;

        // 解析类型
        if (tokenIndex < tokens.size())
        {
            auto type = std::make_unique<BasicTypeNode>(tokens[tokenIndex].text);
            constDecl->bType = std::move(type);
            tokenIndex++;
        }

        // 解析常量定义列表
        do
        {
            if (tokenIndex < tokens.size() && tokens[tokenIndex].type == "IDN")
            {
                string constName = tokens[tokenIndex].text;
                auto constDef = std::make_unique<ConstDefNode>(constName);
                tokenIndex++;

                // 跳过 =
                if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == "=")
                {
                    tokenIndex++;

                    // 解析初始值
                    auto initVal = std::make_unique<ConstInitValNode>();
                    auto expr = parseExpression();
                    if (expr)
                    {
                        initVal->expr = std::move(expr);
                    }
                    constDef->initVal = std::move(initVal);
                }

                constDecl->addChild(std::move(constDef));
            }

            // 检查逗号
            if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == ",")
            {
                tokenIndex++;
            }
            else
            {
                break;
            }
        } while (tokenIndex < tokens.size());

        // 跳过分号
        if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == ";")
        {
            tokenIndex++;
        }

        return std::move(constDecl);
    }

    std::unique_ptr<ASTNode> parseVarDecl()
    {
        // float c;
        auto varDecl = std::make_unique<VarDeclNode>();

        // 解析类型
        if (tokenIndex < tokens.size())
        {
            auto type = std::make_unique<BasicTypeNode>(tokens[tokenIndex].text);
            varDecl->bType = std::move(type);
            tokenIndex++;
        }

        // 解析变量定义列表
        do
        {
            if (tokenIndex < tokens.size() && tokens[tokenIndex].type == "IDN")
            {
                string varName = tokens[tokenIndex].text;
                auto varDef = std::make_unique<VarDefNode>(varName);
                tokenIndex++;

                // 检查初始化
                if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == "=")
                {
                    tokenIndex++;

                    auto initVal = std::make_unique<InitValNode>();
                    auto expr = parseExpression();
                    if (expr)
                    {
                        initVal->expr = std::move(expr);
                    }
                    varDef->initVal = std::move(initVal);
                }

                varDecl->addChild(std::move(varDef));
            }

            // 检查逗号
            if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == ",")
            {
                tokenIndex++;
            }
            else
            {
                break;
            }
        } while (tokenIndex < tokens.size());

        // 跳过分号
        if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == ";")
        {
            tokenIndex++;
        }

        return std::move(varDecl);
    }

    std::unique_ptr<ASTNode> parseFuncDef()
    {
        // float func3(int a1, int A2) { ... }

        // 获取返回类型
        auto returnType = std::make_unique<BasicTypeNode>(tokens[tokenIndex].text);
        tokenIndex++;

        // 获取函数名
        string funcName = "";
        if (tokenIndex < tokens.size())
        {
            if (tokens[tokenIndex].type == "IDN")
            {
                funcName = tokens[tokenIndex].text;
            }
            else if (tokens[tokenIndex].text == "main")
            {
                funcName = "main";
            }
            tokenIndex++;
        }

        auto funcDef = std::make_unique<FuncDefNode>(funcName);
        funcDef->returnType = std::move(returnType);

        // 跳过左括号
        if (tokenIndex < tokens.size())
        {
            tokenIndex++;
        }

        // 解析参数列表
        auto params = std::make_unique<CompUnitNode>(); // 临时用CompUnit存储参数
        while (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) != ")")
        {
            if (tokens[tokenIndex].type == "KW" &&
                (tokens[tokenIndex].text == "int" || tokens[tokenIndex].text == "float"))
            {

                auto paramType = std::make_unique<BasicTypeNode>(tokens[tokenIndex].text);
                tokenIndex++;

                if (tokenIndex < tokens.size() && tokens[tokenIndex].type == "IDN")
                {
                    string paramName = tokens[tokenIndex].text;
                    auto param = std::make_unique<FuncFParamNode>(paramName);
                    param->paramType = std::move(paramType);
                    params->addChild(std::move(param));
                    tokenIndex++;
                }
            }

            // 跳过逗号
            if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == ",")
            {
                tokenIndex++;
            }
        }

        if (!params->children.empty())
        {
            funcDef->params = std::move(params);
        }

        // 跳过右括号
        if (tokenIndex < tokens.size())
        {
            tokenIndex++;
        }

        // 解析函数体
        auto body = parseBlock();
        if (body)
        {
            funcDef->body = std::move(body);
        }

        return std::move(funcDef);
    }

    std::unique_ptr<ASTNode> parseBlock()
    {
        auto block = std::make_unique<BlockNode>();

        // 跳过左大括号
        if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == "{")
        {
            tokenIndex++;
        }

        // 解析语句列表
        while (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) != "}")
        {
            auto stmt = parseStatement();
            if (stmt)
            {
                block->addChild(std::move(stmt));
            }
            else
            {
                // 跳过无法识别的token
                tokenIndex++;
            }
        }

        // 跳过右大括号
        if (tokenIndex < tokens.size())
        {
            tokenIndex++;
        }

        return std::move(block);
    }

    std::unique_ptr<ASTNode> parseStatement()
    {
        if (tokenIndex >= tokens.size())
            return nullptr;

        string terminal = mapTokenTypeToTerminal(tokens[tokenIndex]);

        // 检查是否是声明语句（在函数体内）
        if (terminal == "const" || terminal == "int" || terminal == "float")
        {
            if (terminal == "const")
            {
                return parseConstDecl();
            }
            else
            {
                return parseVarDecl();
            }
        }
        else if (terminal == "return")
        {
            auto returnStmt = std::make_unique<ReturnStmtNode>();
            tokenIndex++; // 跳过return

            // 解析返回表达式（如果有）
            if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) != ";")
            {
                auto expr = parseExpression();
                if (expr)
                {
                    returnStmt->expr = std::move(expr);
                }
            }

            // 跳过分号
            if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == ";")
            {
                tokenIndex++;
            }

            return std::move(returnStmt);
        }
        else if (tokens[tokenIndex].type == "IDN")
        {
            // 函数调用语句  ***********************************开始***************************

            if (tokenIndex + 1 < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex + 1]) == "(")
            {
                string funcName = tokens[tokenIndex].text;
                tokenIndex += 2;

                auto CallExprStmt = std::make_unique<CallExprNode>(funcName);
                CallExprStmt->funcName = funcName;

                // 参数是expr(，expr)*
                auto expr = parseExpression();
                if (expr)
                {
                    CallExprStmt->args.push_back(std::move(expr));
                }

                while (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == ",")
                {
                    tokenIndex++;
                    auto expr = parseExpression();
                    if (expr)
                    {
                        CallExprStmt->args.push_back(std::move(expr));
                    }
                }

                // 跳过括号
                if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == ")")
                {
                    tokenIndex++;
                }
                // 跳过分号
                if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == ";")
                {
                    tokenIndex++;
                }

                return std::move(CallExprStmt);
            }
            //********************************调用结束*********************************************
            else
            {
                // 赋值语句
                string varName = tokens[tokenIndex].text;
                tokenIndex++;

                if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == "=")
                {
                    auto assignStmt = std::make_unique<AssignStmtNode>();
                    assignStmt->lval = std::make_unique<IdentExprNode>(varName);
                    tokenIndex++; // 跳过 =

                    auto expr = parseExpression();
                    if (expr)
                    {
                        assignStmt->expr = std::move(expr);
                    }

                    // 跳过分号
                    if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == ";")
                    {
                        tokenIndex++;
                    }

                    return std::move(assignStmt);
                }
            }
        }
        else if (terminal == "if")
        {
            auto ifStmt = std::make_unique<IfStmtNode>();
            tokenIndex++; // 跳过if

            // 跳过左括号
            if (tokenIndex < tokens.size())
                tokenIndex++;

            // 解析条件表达式
            auto condition = parseExpression();
            if (condition)
            {
                ifStmt->condition = std::move(condition);
            }

            // 跳过右括号
            if (tokenIndex < tokens.size())
                tokenIndex++;

            // 解析then语句（可能是语句块）
            auto thenStmt = parseStatement();
            if (thenStmt)
            {
                ifStmt->thenStmt = std::move(thenStmt);
            }

            // 检查else
            if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == "else")
            {
                tokenIndex++; // 跳过else
                auto elseStmt = parseStatement();
                if (elseStmt)
                {
                    ifStmt->elseStmt = std::move(elseStmt);
                }
            }

            return std::move(ifStmt);
        }
        else if (terminal == "{")
        {
            // 语句块
            return parseBlock();
        }
        else if (terminal == ";")
        {
            // 空语句
            tokenIndex++;
            return nullptr; // 不创建节点
        }

        return nullptr;
    }

    std::unique_ptr<ASTNode> parseExpression()
    {
        return parseLogicalOrExpression();
    }

    std::unique_ptr<ASTNode> parseLogicalOrExpression()
    {
        auto left = parseLogicalAndExpression();

        while (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == "||")
        {
            string op = tokens[tokenIndex].text;
            tokenIndex++;
            auto right = parseLogicalAndExpression();

            auto binaryExpr = std::make_unique<BinaryExprNode>(op);
            binaryExpr->left = std::move(left);
            binaryExpr->right = std::move(right);
            left = std::move(binaryExpr);
        }

        return left;
    }

    std::unique_ptr<ASTNode> parseLogicalAndExpression()
    {
        auto left = parseEqualityExpression();

        while (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == "&&")
        {
            string op = tokens[tokenIndex].text;
            tokenIndex++;
            auto right = parseEqualityExpression();

            auto binaryExpr = std::make_unique<BinaryExprNode>(op);
            binaryExpr->left = std::move(left);
            binaryExpr->right = std::move(right);
            left = std::move(binaryExpr);
        }

        return left;
    }

    std::unique_ptr<ASTNode> parseEqualityExpression()
    {
        auto left = parseRelationalExpression();

        while (tokenIndex < tokens.size() &&
               (mapTokenTypeToTerminal(tokens[tokenIndex]) == "==" ||
                mapTokenTypeToTerminal(tokens[tokenIndex]) == "!="))
        {
            string op = tokens[tokenIndex].text;
            tokenIndex++;
            auto right = parseRelationalExpression();

            auto binaryExpr = std::make_unique<BinaryExprNode>(op);
            binaryExpr->left = std::move(left);
            binaryExpr->right = std::move(right);
            left = std::move(binaryExpr);
        }

        return left;
    }

    std::unique_ptr<ASTNode> parseRelationalExpression()
    {
        auto left = parseAdditiveExpression();

        while (tokenIndex < tokens.size() &&
               (mapTokenTypeToTerminal(tokens[tokenIndex]) == "<" ||
                mapTokenTypeToTerminal(tokens[tokenIndex]) == ">" ||
                mapTokenTypeToTerminal(tokens[tokenIndex]) == "<=" ||
                mapTokenTypeToTerminal(tokens[tokenIndex]) == ">="))
        {
            string op = tokens[tokenIndex].text;
            tokenIndex++;
            auto right = parseAdditiveExpression();

            auto binaryExpr = std::make_unique<BinaryExprNode>(op);
            binaryExpr->left = std::move(left);
            binaryExpr->right = std::move(right);
            left = std::move(binaryExpr);
        }

        return left;
    }

    std::unique_ptr<ASTNode> parseAdditiveExpression()
    {
        auto left = parseMultiplicativeExpression();

        while (tokenIndex < tokens.size() &&
               (mapTokenTypeToTerminal(tokens[tokenIndex]) == "+" ||
                mapTokenTypeToTerminal(tokens[tokenIndex]) == "-"))
        {
            string op = tokens[tokenIndex].text;
            tokenIndex++;
            auto right = parseMultiplicativeExpression();

            auto binaryExpr = std::make_unique<BinaryExprNode>(op);
            binaryExpr->left = std::move(left);
            binaryExpr->right = std::move(right);
            left = std::move(binaryExpr);
        }

        return left;
    }

    std::unique_ptr<ASTNode> parseMultiplicativeExpression()
    {
        auto left = parsePrimaryExpression();

        while (tokenIndex < tokens.size() &&
               (mapTokenTypeToTerminal(tokens[tokenIndex]) == "*" ||
                mapTokenTypeToTerminal(tokens[tokenIndex]) == "/" ||
                mapTokenTypeToTerminal(tokens[tokenIndex]) == "%"))
        {
            string op = tokens[tokenIndex].text;
            tokenIndex++;
            auto right = parsePrimaryExpression();

            auto binaryExpr = std::make_unique<BinaryExprNode>(op);
            binaryExpr->left = std::move(left);
            binaryExpr->right = std::move(right);
            left = std::move(binaryExpr);
        }

        return left;
    }

    std::unique_ptr<ASTNode> parsePrimaryExpression()
    {
        if (tokenIndex >= tokens.size())
            return nullptr;

        if (tokens[tokenIndex].type == "INT")
        {
            auto literal = std::make_unique<LiteralExprNode>(tokens[tokenIndex].text, "int");
            tokenIndex++;
            return std::move(literal);
        }
        else if (tokens[tokenIndex].type == "FLOAT")
        {
            auto literal = std::make_unique<LiteralExprNode>(tokens[tokenIndex].text, "float");
            tokenIndex++;
            return std::move(literal);
        }
        else if (tokens[tokenIndex].type == "IDN")
        {
            // 函数调用表达式  ********************开始***********************
            if (tokenIndex + 1 < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex + 1]) == "(")
            {
                string funcName = tokens[tokenIndex].text;
                tokenIndex += 2;

                auto CallExprStmt = std::make_unique<CallExprNode>(funcName);
                CallExprStmt->funcName = funcName;

                // 参数是expr(，expr)*
                auto expr = parseExpression();
                if (expr)
                {
                    CallExprStmt->args.push_back(std::move(expr));
                }

                while (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == ",")
                {
                    tokenIndex++;
                    auto expr = parseExpression();
                    if (expr)
                    {
                        CallExprStmt->args.push_back(std::move(expr));
                    }
                }

                // 跳过括号
                if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == ")")
                {
                    tokenIndex++;
                }
                // 跳过分号
                if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == ";")
                {
                    tokenIndex++;
                }
                return std::move(CallExprStmt);
            }
            // 函数调用表达式  ********************结束***********************
            else
            {
                auto ident = std::make_unique<IdentExprNode>(tokens[tokenIndex].text);
                tokenIndex++;
                return std::move(ident);
            }
        }
        else if (mapTokenTypeToTerminal(tokens[tokenIndex]) == "(")
        {
            tokenIndex++; // 跳过左括号
            auto expr = parseExpression();
            if (tokenIndex < tokens.size() && mapTokenTypeToTerminal(tokens[tokenIndex]) == ")")
            {
                tokenIndex++; // 跳过右括号
            }
            return expr;
        }

        return nullptr;
    }
};

int main()
{
    // 读取词法分析器token输出
    vector<Symbol> tokens = loadTokens("../parser/token.tsv");
    if (tokens.empty())
    {
        cerr << "无输入token，程序退出" << endl;
        return 1;
    }

    cout << "开始构建AST（使用递归下降分析器）..." << endl;

    // 使用递归下降AST构建器
    RecursiveDescentASTBuilder astBuilder(tokens);
    auto ast = astBuilder.buildAST();

    if (ast)
    {
        cout << "\n=== AST 构建成功 ===" << endl;
        cout << "使用内置printTree方法:" << endl;
        ast->printTree();

        cout << "\n使用ASTPrinter:" << endl;
        ASTPrinter printer;
        printer.print(*ast);

        ofstream fout("../irgenerator/test/test.txt");
        if (fout)
        {
            // 保存原 cout 缓冲区
            streambuf *oldCoutBuf = cout.rdbuf();
            cout.rdbuf(fout.rdbuf()); // 重定向 cout 到文件

            // 重复输出一次到文件
            printer.print(*ast);

            // 恢复 cout
            cout.rdbuf(oldCoutBuf);
            fout.close();

            cout << "\n语法树已写入 test.txt" << endl;
        }
        else
        {
            cerr << "无法创建 test.txt 文件" << endl;
        }
    }
    else
    {
        cout << "AST构建失败" << endl;
    }

    return 0;
}